# -*- coding: UTF-8 -*-
"""
@Date    ：2025/9/15 11:15 
@Author  ：Liu Yuezhao
@Project ：bert 
@File    ：pre_time_bert_model.py
@IDE     ：PyCharm 
"""
import torch
from torch import nn
from transformers import BertConfig, BertForMaskedLM
from src.tools.utils import load_config
from .time2vec import TripleTime2Vec
from typing import Optional, Union

# 加载配置
yaml_config = load_config("../../config.yaml")
device = yaml_config["config"]["device"]

# BERT 配置
event_seq_bert_config = BertConfig(
    vocab_size=yaml_config["pretrain_model_config"]["vocab_size"],
    # vocab_size=10000,

    hidden_size=yaml_config["pretrain_model_config"]["hidden_size"],
    num_hidden_layers=yaml_config["pretrain_model_config"]["num_hidden_layers"],
    num_attention_heads=yaml_config["pretrain_model_config"]["num_attention_heads"],
    intermediate_size=yaml_config["pretrain_model_config"]["intermediate_size"],
    max_position_embeddings=yaml_config["pretrain_model_config"]["max_position_embeddings"],
    hidden_act=yaml_config["pretrain_model_config"]["hidden_act"],
    hidden_dropout_prob=yaml_config["pretrain_model_config"]["hidden_dropout_prob"],
    attention_probs_dropout_prob=yaml_config["pretrain_model_config"]["attention_probs_dropout_prob"],
    pad_token_id=0,
    type_vocab_size=2,
)


class TimeBertEmbedding(nn.Module):
    def __init__(self, config, time2vec_dim=768, time_activation="cos"):
        super().__init__()
        # bert 默认的 3 个embedding
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        # self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        # self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # 时间属性的 3 个embedding
        self.triple_time2vec = TripleTime2Vec(
            embed_dim=time2vec_dim,
            activation=time_activation
        )

        self.projection = nn.Linear(3 * time2vec_dim, config.hidden_size)

        # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        # self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        # self.register_buffer(
        #     "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        # )
        # self.register_buffer(
        #     "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
        # )

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values_length: int = 0,
        interval: torch.Tensor = None,
        same_event_interval: torch.Tensor = None,
        age: torch.Tensor = None,
    ) -> torch.Tensor:
        # if input_ids is not None:
        #     input_shape = input_ids.size()
        # else:
        #     input_shape = inputs_embeds.size()[:-1]
        #
        # seq_length = input_shape[1]
        #
        # if position_ids is None:
        #     position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
        #
        # if token_type_ids is None:
        #     if hasattr(self, "token_type_ids"):
        #         buffered_token_type_ids = self.token_type_ids[:, :seq_length]
        #         buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
        #         token_type_ids = buffered_token_type_ids_expanded
        #     else:
        #         token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        # 添加时间信息embedding
        interval_embed, same_event_embed, age_embed = self.triple_time2vec(
            interval=interval,
            same_event_interval=same_event_interval,
            age=age
        )

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)

        # 拼接token + time2vec * 2
        concat_embeds = torch.cat([inputs_embeds, interval_embed, same_event_embed], dim=-1)
        # 维度变换回来
        concat_embeds = self.projection(concat_embeds)

        # token_type_embeddings = self.token_type_embeddings(token_type_ids)

        # embeddings = concat_embeds + token_type_embeddings
        # if self.position_embedding_type == "absolute":
        #     position_embeddings = self.position_embeddings(position_ids)
        #     embeddings += position_embeddings

        embeddings = concat_embeds + age_embed

        # embeddings = self.LayerNorm(embeddings)
        # embeddings = self.dropout(embeddings)
        return embeddings


class SeqBertPretrain(nn.Module):
    def __init__(
        self,
        config: BertConfig,
        time2vec_dim = 768,
        time_activation = 'cos'
    ):
        """
        支持三时间信号的 BERT 预训练模型
        """
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size

        # embedding
        self.emb = TimeBertEmbedding(config, time2vec_dim=time2vec_dim, time_activation=time_activation)

        # BERT MLM 主干
        self.bert_mlm = BertForMaskedLM(config=config)

        # 初始化
        self.apply(self._init_weights)

    def _init_weights(self, module):
        """初始化参数"""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if hasattr(module, 'bias') and module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(
            self,
            input_ids: torch.Tensor = None,
            attention_mask: torch.Tensor = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            labels: torch.Tensor = None,
            interval: torch.Tensor = None,
            same_event_interval: torch.Tensor = None,
            age: torch.Tensor = None,
            inputs_embeds: torch.Tensor = None,
            past_key_values_length: int = 0,
    ):
        embeds = self.emb(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
            interval=interval,
            same_event_interval=same_event_interval,
            age=age,
        )

        outputs = self.bert_mlm(
            inputs_embeds=embeds,
            attention_mask=attention_mask,
            labels=labels,
            return_dict=True
        )
        return outputs


# 初始化模型
pre_bert_model = SeqBertPretrain(
    config=event_seq_bert_config,
    time2vec_dim=yaml_config["pretrain_model_config"]["hidden_size"],
    time_activation='cos'
).to(device)